package ner;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Set;

import org.xml.sax.Attributes;
import org.xml.sax.InputSource;
import org.xml.sax.SAXException;
import org.xml.sax.XMLReader;
import org.xml.sax.helpers.DefaultHandler;
import org.xml.sax.helpers.XMLReaderFactory;

public class Testing extends DefaultHandler {

	private String mostUsedNETag;
	private HashMap<String, String> vocabMap_;
	private Training training_; 
	//	private HashMap<String, HashMap<String, Integer>> wordGivenTagWord;
	private String outputFile;
	BufferedWriter outputWriter;
	private String currTag = Training.START_TAG;
	private String prevTag = Training.START_TAG;
	private String prevWord = Training.START_WORD;
	private ArrayList<String> words = new ArrayList<String>(226000);
	private int unk = 0;

	public Testing(Training training) {
		this.mostUsedNETag = training.getMostUsedNETag();
		this.vocabMap_ = training.getWordMap();
		this.training_ = training;
	}

	/**
	 * @param args
	 */
	public static void main(String[] args) {

		Training training = new Training();

		String inputFile = "data/ne_files/ne.muc6.training.keys";
		String outFile = "data/ne_files/ne.muc6.testing.output";
		String testFile = "data/ne_files/ne.muc6.test.texts";

		for(int i = 0; i < args.length; i = i+2) {
			if(args[i].equalsIgnoreCase("-i"))
				inputFile = args[i+1];
			if(args[i].equalsIgnoreCase("-o"))
				outFile = args[i+1];
			if(args[i].equalsIgnoreCase("-t"))
				testFile = args[i+1];
		}
		

		// intermediateTestFile.
		String keyFile = "data/ne_files/ne.muc6.test.keys";

		// intermediateTestFile.
		String keyoutputFile = "data/ne_files/ne.muc6.test.keys.output";

		// intermediateTrainingFile.
		String interTrainFile = "data/ne_files/ne.muc6.training.intermediate";
		
		// intermediateTestFile.
		String interTestFile = "data/ne_files/ne.muc6.testing.intermediate";
		
		training.train(inputFile, interTrainFile);

		System.out.println("TRAINING DONE");
		Testing test = new Testing(training);
		test.test(testFile, interTestFile, outFile);

		KeyGenerator keyGenerator = new KeyGenerator();
		keyGenerator.generateKeys(keyFile, keyoutputFile, interTestFile);

		Scorer scorer = new Scorer();
		scorer.calculateScore(outFile, keyoutputFile);
	}

	private void test(String inputFile, String interTestFile, String outputFile) {

		try {
			this.outputFile = outputFile;
			this.outputWriter = new BufferedWriter(new FileWriter(this.outputFile));

			/* The new file will be called raw.txt */
			FileWriter writer = new FileWriter(interTestFile);
			FileReader reader = new FileReader(inputFile);

			XMLReader xmlr = XMLReaderFactory.createXMLReader();

			xmlr.setContentHandler(this);
			xmlr.setErrorHandler(this);

			BufferedReader br = new BufferedReader(reader);

			try {
				format(br, writer);
				reader.close();
			}
			catch (IOException ex) {
				throw new RuntimeException(ex);
			}

			reader = new FileReader(interTestFile);

			// Parse the incoming XML file.
			xmlr.parse(new InputSource(reader));
		}
		catch (IOException ex) {
			throw new RuntimeException(ex);
		}
		catch (SAXException e) {
			throw new RuntimeException(e);
		}
	}
	public void format(BufferedReader br, FileWriter out)
	throws IOException
	{
		String line;
		try {
			out.write("<ROOT>" + "\n");
			while ((line = br.readLine()) != null) {
				line = line.replaceAll("&", "@amp;");
				out.write(line.trim() + "\n");
			}
			out.write("</ROOT>" + "\n");
		}
		catch (IOException ex) {
			System.err.println(ex);
		}

		out.close();
		br.close();
	}
	/*
	 * Grabs the opening SGML tag. 
	 */
	@Override
	public void startElement(String uri, String name, String qName, Attributes atts)
	{
	}
	/*
	 * Grabs the closing tag. 
	 */
	@Override
	public void endElement(String uri, String name, String qName)
	{
	}
	/*
	 * This prints out all the text between the tags we care about to 
	 * a file. 
	 *
	 */
	@Override
	public void characters(char ch[], int start, int length)
	{
		//		try {
		String text = new String(ch, start, length);

		String uTxt = unescapeText(text);
		uTxt = uTxt.replaceAll("'s", "\n's").replaceAll(",", "\n,").replaceAll("`", "`\n").replaceAll("[(]", "(\n");

		String[] strings = uTxt.split(" |[.]|\\n");

		if (strings!=null) {
			for (String string : strings) {
				if (string!=null && !string.isEmpty()) {
					words.add(string);
				}
			}
		}
	}
	
	
	public void viterbiHMM () {
		try {
			Set<String> tagSet = training_.getTagSet();
			String[] tagsArray = tagSet.toArray(new String[2]);
			double[][] score = new double[tagsArray.length][words.size()];
			int[][] bptr = new int[tagsArray.length][words.size()];

			for(int t = 0; t < tagsArray.length; t++) {
				String key1 = combine(tagsArray[t], Training.START_WORD);
				double term1 = training_.probWGTW(words.get(0), key1);
				String key2 = combine(Training.START_TAG,Training.START_WORD);	
				double term2 = training_.probTGTW(tagsArray[t], key2);
				score[t][0] = Math.log( term1 * term2);
				bptr[t][0] = 0;
			}

			for(int w = 1; w < words.size(); w++) {
				for(int t = 0; t < tagsArray.length; t++) {
					double max = Double.NEGATIVE_INFINITY;
					int bestPrev = -1;
					for(int p = 0; p < tagsArray.length; p++) {
						String key = combine(tagsArray[p], words.get(w-1));	
						double term2 = training_.probTGTW(tagsArray[t], key);
						double tmpScore = score[p][w-1] + Math.log(term2);
						if(tmpScore > max) {
							max = tmpScore;
							bestPrev = p;
						}
					}
					score[t][w] = max + Math.log(training_.probWGTW(words.get(w), 
							combine(tagsArray[t], words.get(w-1))));
					bptr[t][w] = bestPrev;
				}
			}
			
			double max = Double.NEGATIVE_INFINITY;
			int bestEndTag = -1;
			for(int t = 0; t < tagsArray.length; t++) {
				if(score[t][words.size()-1] > max) {
					max = score[t][words.size()-1];
					bestEndTag = t;
				}
			}

			int[] tags = new int[words.size()];
			tags[words.size()-1] = bestEndTag;
			for(int w = words.size()-2; w>=0; w--) {
				tags[w] = bptr[tags[w+1]][w+1];
			}

			for(int w = 0; w < words.size(); w++) {
				String word = words.get(w);
				if(!vocabMap_.containsKey(word)) {
					if(Character.isUpperCase(word.charAt(0))) {
						outputWriter.write(mostUsedNETag + " " + words.get(w) + "\n");
						unk ++;
						continue;
					}
				}
				outputWriter.write(tagsArray[tags[w]] + " " + words.get(w) + "\n");
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
	}
	private String combine(String tag, String word) {
		return tag+ Training.SEPRTR + word;
	}

	public String unescapeText(String text)
	{
		text = text.replaceAll("@amp;", "&");
		text = text.replaceAll("&MD;", "-");
		text = text.replaceAll("&AMP;", "&");
		text = text.replaceAll("&LR;", "");

		text = text.replaceAll("``|''", "\"");
		return text;
	}
	@Override
	public void endDocument() throws SAXException {
		try {
			viterbiHMM();
			System.out.println("Unknown " + unk);
			outputWriter.flush();
			outputWriter.close();
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

}
